from datetime import timedelta, datetime


class TimeSlot(object):
    '''给定时间段内月份or日期'''

    def __init__(self, start, end):
        self.start = start
        self.end = end

    def gen_dates(self, b_date, days):
        day = timedelta(days=1)
        # print(day)
        for i in range(days):
            # print(b_date + day*i)
            yield b_date + day * i

    def get_date_list(self):  # end_date=None
        """
        获取日期列表
        :param start: 开始日期
        :param end: 结束日期
        :return:
        """
        if self.start is not None:
            start = datetime.strptime(self.start, "%Y-%m-%d")
        else:
            start = datetime.now()
        if self.end is None:
            end = datetime.now()
        else:
            end = datetime.strptime(self.end, "%Y-%m-%d")
        data = []
        for d in self.gen_dates(start, ((end - start).days + 1)):  # 29 + 1
            # print(d)   # datetime.datetime  类型
            data.append(d.strftime("%Y-%m-%d"))
        return data

    def get_month_list(self):
        '''获取时间段内月份'''
        dates = self.get_date_list()
        months = []
        for i in dates:
            if i[:7] not in months:
                months.append(i[:7])
        return months


if __name__ == '__main__':
    time_slot = TimeSlot('2020-01-01', '2020-12-31')
    date_list = time_slot.get_date_list()
    print(date_list)
    month_list = time_slot.get_month_list()
    print(month_list)
